import random

from torchvision import transforms
from pathlib import Path
import numpy as np
import torch

from PIL import Image
from torch.utils.data.dataset import Dataset
from torchvision.transforms import transforms

from .utils import cvtColor, preprocess_input, resize_image



# 自定义数据集所需要的类
class MyDataset(Dataset):
    def __init__(self, dir_path):
        self.dir_path = dir_path  # 数据集根目录
        self.transform = transforms.ToTensor()
        self.images, self.labels = self.get_images_and_labels(self.dir_path)

    def __len__(self):
        # 返回数据集的数据数量
        return len(self.images)

    def __getitem__(self, index):
        # 传入index 计算公式为people_index*5+img_index (0-499)
        img_path = self.images[index]
        target = self.labels[index//8]
        img_pil = Image.open(img_path)   # 使用PIL格式先
        img_tensor = self.transform(img_pil)  # 转换成tensor格式
        # sample = {'image': img_tensor, 'label': label}
        return img_tensor, target

    # 自定义数据集所需要的方法
    def get_images_and_labels(self, dir_path):
        """
        从图像数据集的根目录dir_path下获取所有类别的图像名列表和对应的标签名列表
        :param dir_path: 图像数据集的根目录
        :return: images_list, labels_list
        """
        dir_path = Path(dir_path)
        classes = {}  # 类别名列表
        # 三级文件
        for category in dir_path.iterdir():
            if category.is_dir():
                classes[category.name] = []
                son_categories = category.iterdir()
                for son_category in son_categories:
                    classes[category.name].append(son_category.name)  # 把各个文件夹名字写入到classes 如0、1、2...498、499

        images_list = []  # 文件名列表
        labels_list = []  # 标签列表
        count = 0
        for key, value in classes.items():
            for index, son_value in enumerate(value):
                class_path = dir_path / key / son_value
                if not class_path.is_dir():
                    continue
                for img_path in class_path.glob('*.bmp'):
                    images_list.append(str(img_path))
                labels_list.append(count)
                count = count + 1

        return images_list, labels_list


# facenet数据集
class FacenetDataset(Dataset):
    def __init__(self, input_shape, lines, num_classes, random):
        self.input_shape = input_shape
        self.lines = lines
        self.length = len(lines)
        self.num_classes = num_classes
        self.random = random

        # ------------------------------------#
        #   路径和标签
        # ------------------------------------#
        self.paths = []
        self.labels = []

        self.load_dataset()

    def __len__(self):
        return self.length

    def __getitem__(self, index):
        # ------------------------------------#
        #   创建全为零的矩阵
        # ------------------------------------#
        images = np.zeros((3, 3, self.input_shape[0], self.input_shape[1]))
        labels = np.zeros(3)

        # ------------------------------#
        #   先获得两张同一个人的人脸
        #   用来作为anchor和positive
        # ------------------------------#
        c = random.randint(0, self.num_classes - 1)
        selected_path = self.paths[self.labels[:] == c]
        while len(selected_path) < 2:
            c = random.randint(0, self.num_classes - 1)
            selected_path = self.paths[self.labels[:] == c]

        # ------------------------------------#
        #   随机选择两张
        # ------------------------------------#
        image_indexes = np.random.choice(range(0, len(selected_path)), 2)
        # ------------------------------------#
        #   打开图片并放入矩阵
        # ------------------------------------#
        image = cvtColor(Image.open(selected_path[image_indexes[0]]))
        # ------------------------------------------#
        #   翻转图像
        # ------------------------------------------#
        if self.rand() < .5 and self.random:
            image = image.transpose(Image.FLIP_LEFT_RIGHT)
        image = resize_image(image, [self.input_shape[1], self.input_shape[0]], letterbox_image=True)
        image = preprocess_input(np.array(image, dtype='float32'))
        image = np.transpose(image, [2, 0, 1])
        images[0, :, :, :] = image
        labels[0] = c

        image = cvtColor(Image.open(selected_path[image_indexes[1]]))
        # ------------------------------------------#
        #   翻转图像
        # ------------------------------------------#
        if self.rand() < .5 and self.random:
            image = image.transpose(Image.FLIP_LEFT_RIGHT)
        image = resize_image(image, [self.input_shape[1], self.input_shape[0]], letterbox_image=True)
        image = preprocess_input(np.array(image, dtype='float32'))
        image = np.transpose(image, [2, 0, 1])
        images[1, :, :, :] = image
        labels[1] = c

        # ------------------------------#
        #   取出另外一个人的人脸
        # ------------------------------#
        different_c = list(range(self.num_classes))
        different_c.pop(c)
        different_c_index = np.random.choice(range(0, self.num_classes - 1), 1)
        current_c = different_c[different_c_index[0]]
        selected_path = self.paths[self.labels == current_c]
        while len(selected_path) < 1:
            different_c_index = np.random.choice(range(0, self.num_classes - 1), 1)
            current_c = different_c[different_c_index[0]]
            selected_path = self.paths[self.labels == current_c]

        # ------------------------------#
        #   随机选择一张
        # ------------------------------#
        image_indexes = np.random.choice(range(0, len(selected_path)), 1)
        image = cvtColor(Image.open(selected_path[image_indexes[0]]))
        # ------------------------------------------#
        #   翻转图像
        # ------------------------------------------#
        if self.rand() < .5 and self.random:
            image = image.transpose(Image.FLIP_LEFT_RIGHT)
        image = resize_image(image, [self.input_shape[1], self.input_shape[0]], letterbox_image=True)
        image = preprocess_input(np.array(image, dtype='float32'))
        image = np.transpose(image, [2, 0, 1])
        images[2, :, :, :] = image
        labels[2] = current_c

        return images, labels

    def rand(self, a=0, b=1):
        return np.random.rand() * (b - a) + a

    def load_dataset(self):
        for path in self.lines:
            path_split = path.split(";")
            self.paths.append(path_split[1].split()[0])
            self.labels.append(int(path_split[0]))
        self.paths = np.array(self.paths, dtype=np.object)
        self.labels = np.array(self.labels)



# DataLoader中collate_fn使用  是在使用triplet loss 时使用
# 将三张图打包成batch
def dataset_collate(batch):
    images = []
    labels = []
    for img, label in batch:
        images.append(img)
        labels.append(label)

    images1 = np.array(images)[:, 0, :, :, :]
    images2 = np.array(images)[:, 1, :, :, :]
    images3 = np.array(images)[:, 2, :, :, :]
    images = np.concatenate([images1, images2, images3], 0)

    labels1 = np.array(labels)[:, 0]
    labels2 = np.array(labels)[:, 1]
    labels3 = np.array(labels)[:, 2]
    labels = np.concatenate([labels1, labels2, labels3], 0)

    images = torch.from_numpy(np.array(images)).type(torch.FloatTensor)
    labels = torch.from_numpy(np.array(labels)).long()
    return images, labels


# SiameseNet孪生网络数据集
class SiameseNetworkDataset(Dataset):
    def __init__(self, input_shape, imginfo, random):
        self.input_shape = input_shape
        self.imginfo = imginfo
        self.length = len(imginfo)
        self.random = random

        # ------------------------------------#
        #   标签和两张图的路径
        # ------------------------------------#
        self.dataDic = dict()
        self.load_dataset()  # 加载数据 以更新self.dataDic

    def __len__(self):
        return self.length

    def __getitem__(self, index):
        # ------------------------------------#
        #   随机选择一组(两张图片)
        # ------------------------------------#
        c = random.randint(0, self.length - 1)
        # ------------------------------------#
        #   打开图片并放入矩阵
        # ------------------------------------#
        imageA = cvtColor(Image.open(self.dataDic[str(c)][1]))
        imageB = cvtColor(Image.open(self.dataDic[str(c)][2]))
        # ------------------------------------------#
        #   翻转图像 为了增加模型的泛化能力
        # ------------------------------------------#
        if self.rand() < .5 and self.random:
            imageA = imageA.transpose(Image.FLIP_LEFT_RIGHT)
        elif self.rand() > .5 and self.random:
            imageB = imageB.transpose(Image.FLIP_LEFT_RIGHT)

        # 一系列的处理
        imageA = resize_image(imageA, [self.input_shape[1], self.input_shape[0]], letterbox_image=True)       # resize成(160,160)
        imageB = resize_image(imageB, [self.input_shape[1], self.input_shape[0]], letterbox_image=True)
        imageA = preprocess_input(np.array(imageA, dtype='float32'))                                          # 像素点颜色/255
        imageB = preprocess_input(np.array(imageB, dtype='float32'))
        imageA = np.transpose(imageA, [2, 0, 1])                                                              # 将通道数放在前面索引
        imageB = np.transpose(imageB, [2, 0, 1])

        label = self.dataDic[str(c)][0]

        # 返回
        return imageA, imageB, label

    def rand(self, a=0, b=1):
        return np.random.rand() * (b - a) + a

    # 加载每一条标签和图片path  每一条格式如下"0;0;XXX/XXX/a.jpg;XXX/XXX/b.jpg"
    def load_dataset(self):
        for info in self.imginfo:
            data_info = info.split(";")
            # 格式如下 {'0':[1, XXX/XXX/a.jpg, XXX/XXX/b.jpg]}
            self.dataDic[data_info[0]] = [int(data_info[1]), data_info[2].split()[0], data_info[3].split()[0]]


# 加载LFW数据 用以验证模型的准确性
class LFWDataset(Dataset):
    def __init__(self, pairs_path, image_size):
        self.image_size = image_size
        self.pairs_path = pairs_path
        self.validation_images = self.get_lfw(pairs_path)

    # 获取用以验证模型的lfw数据
    def get_lfw(self, lfw_pair):
        with open(lfw_pair, 'r') as f:
            result = f.readlines()
        return result

    # 处理原始信息
    def parse_info(self, info):
        info_lis = info.split(';')
        issame = info_lis[0]                  # 是否相同标签 0:相同 1:不同
        img1_path = info_lis[1].strip()       # 第一张图片的地址
        img2_path = info_lis[2].strip()       # 第二张图片的地址

        return issame, img1_path, img2_path

    def __getitem__(self, index):
        (issame, path_1, path_2) = self.parse_info(self.validation_images[index])
        image1, image2 = Image.open(path_1), Image.open(path_2)
        image1 = resize_image(image1, [self.image_size[1], self.image_size[0]], letterbox_image=True)
        image2 = resize_image(image2, [self.image_size[1], self.image_size[0]], letterbox_image=True)

        image1, image2 = np.transpose(preprocess_input(np.array(image1, np.float32)), [2, 0, 1]), np.transpose(
            preprocess_input(np.array(image2, np.float32)), [2, 0, 1])

        return image1, image2, issame

    def __len__(self):
        return len(self.validation_images)



